#include <linux/vmalloc.h>
#include <asm/pgtable_types.h>

#include <linux/module.h>

#include "../include/testing_commands.h"

#include "print_test.h"

struct tester_data TESTER;

char *wrex;


int get_SMEP (void) {
  u64 cr4;
  asm volatile("movq %%cr4, %0" : "=r" (cr4));
  return (cr4 >> 17) & 1;
}


void set_comm(void)
{
  //PML4E -> PDPTE
  map   (0, 1, 1);
  flags (0, 1, PF_PRESENT | PF_WRITE);
  
  //PDPTE -> PDE
  map   (1, 1, 2);
  flags (1, 1, PF_PRESENT | PF_WRITE);
  
  //PDE -> PTE
  map   (2, 1, 3);
  flags (2, 1, PF_PRESENT | PF_WRITE);
  
  //PTE -> 4kb page
  map   (3, 1, 4);
  flags (3, 1, PF_PRESENT | PF_WRITE | PF_UNEXECUTABLE);
  
  //PTE -> 4kb page
  map   (5, 1, 6);
  flags (5, 1, PF_PRESENT);
  
  text  ("0 -W> 1 -W> 2 -W> 3 -W XD> 4\n\n"
	 "                    5 -> 6");
  
  addr  ("write 1.1.1.1", CMD_WRITE, vaddr64_5(1,1,1,1,17), 4);
  addr  ("run 1.1.1.1", CMD_RUN, vaddr64_5(1,1,1,1,17), 4);
  addr  ("read 1.1.1.1", CMD_READ, vaddr64_5(1,1,1,1,17), 4);
  
  map   (2, 1, 5);
  text ("remap 2.1 to 5");
  
  addr  ("write 1.1.1.1", CMD_WRITE, vaddr64_5(1,1,1,1,17), 4);
  addr  ("run 1.1.1.1", CMD_RUN, vaddr64_5(1,1,1,1,17), 4);
  addr  ("read 1.1.1.1", CMD_READ, vaddr64_5(1,1,1,1,17), 4);
  
  addr  ("@invlpg", CMD_INV, vaddr64_5(1,1,1,1,17), 4);
  
  addr  ("write 1.1.1.1", CMD_WRITE, vaddr64_5(1,1,1,1,17), 4);
  addr  ("run 1.1.1.1", CMD_RUN, vaddr64_5(1,1,1,1,17), 6);
  addr  ("read 1.1.1.1", CMD_READ, vaddr64_5(1,1,1,1,17), 4);
  
  flags (5, 1, PF_WRITE | PF_UNEXECUTABLE);
  text ("flags 5.1 {W, XD}");
   
  addr  ("write 1.1.1.1", CMD_WRITE, vaddr64_5(1,1,1,1,17), 4);
  addr  ("run 1.1.1.1", CMD_RUN, vaddr64_5(1,1,1,1,17), 6);
  addr  ("read 1.1.1.1", CMD_READ, vaddr64_5(1,1,1,1,17), 4);
  
  addr  ("@invlpg", CMD_INV, vaddr64_5(1,1,1,1,17), 4);
  
  addr  ("write 1.1.1.1", CMD_WRITE, vaddr64_5(1,1,1,1,17), 4);
  addr  ("run 1.1.1.1", CMD_RUN, vaddr64_5(1,1,1,1,17), 6);
  addr  ("read 1.1.1.1", CMD_READ, vaddr64_5(1,1,1,1,17), 4);
  //addr  ("run 1.1.1.1", CMD_RUN, vaddr64_5(1,1,1,1,17), 6);
  
  uflags (5, 1, PF_WRITE | PF_UNEXECUTABLE);
  text ("flags 5.1 {!W, !XD}");
   
  addr  ("write 1.1.1.1", CMD_WRITE, vaddr64_5(1,1,1,1,17), 4);
  addr  ("run 1.1.1.1", CMD_RUN, vaddr64_5(1,1,1,1,17), 6);
  addr  ("read 1.1.1.1", CMD_READ, vaddr64_5(1,1,1,1,17), 4);
  
  flags (1, 1, 1llu << 51);
  text ("flags 5.1 {bit51}");
  
  addr  ("read 1.1.1.1", CMD_READ, vaddr64_5(1,1,1,1,17), 4);
  addr  ("write 1.1.1.1", CMD_WRITE, vaddr64_5(1,1,1,1,17), 4);
  addr  ("read 1.1.1.1", CMD_READ, vaddr64_5(1,1,1,1,17), 4);
  addr  ("run 1.1.1.1", CMD_RUN, vaddr64_5(1,1,1,1,17), 6);
  addr  ("read 1.1.1.1", CMD_READ, vaddr64_5(1,1,1,1,17), 4);
  
  addr  ("@invlpg", CMD_INV, vaddr64_5(1,1,1,1,17), 4);
  
  addr  ("write 1.1.1.1", CMD_WRITE, vaddr64_5(1,1,1,1,17), 4);
  addr  ("run 1.1.1.1", CMD_RUN, vaddr64_5(1,1,1,1,17), 6);
  addr  ("read 1.1.1.1", CMD_READ, vaddr64_5(1,1,1,1,17), 4);
  
  addr  ("flush", CMD_FLUSH, 0, 0);
  addr  ("end", CMD_END, 0, 0);
}




static int my_init(void)
{
  int res;
  wrex = __vmalloc(1, GFP_KERNEL, PAGE_KERNEL_EXEC);
  *wrex = 0xCC;
  
  info("\n=======begin=======\n");
  if (init_tester() != 0)
    return 1;
  set_comm();
  
  info("initialization completed\n");
  
  res = run_test();
  print_test_results(res);
  print_test();
  print_table_n ((struct interrupt_info *)TESTER.results,
								 (struct interrupt_info *)TESTER.result, 3);
  
  free_memory();
  
  debug = get_SMEP();
  info("\ndebug: %llu\n", debug);
  
  info("\nfinish\n");
  
  return 0;
}

static void my_exit(void) {
  info("exit\n");
}

module_init(my_init);
module_exit(my_exit);
